
type 'a directory = 
    Directory of (string * 'a directory list)
  | File of 'a

let make name (l: 'a directory list) : 'a directory = 
    Directory (name,l)

let file (f: 'a) : 'a directory = File f
let of_files name (fl: 'a list) : 'a directory =
    let l' = List.map (fun elt -> File elt) fl in
    Directory (name, l')

let of_dirs name (dl: 'a directory list) : 'a directory =
    Directory (name, dl)
;;
let type_of = function Directory _ -> "directory" | File _ -> "file"

(*let rec map (f: 'a -> 'b) (dir: 'a directory) : 'b directory =*)
let rec map ?(name=None) ~f dir =
    ( match dir with 
      | Directory (name',l) -> 
          let l' = List.map (fun dir -> map ~f dir) l in
	      make (match name with Some s -> s | None -> name') l'
(*          Directory (name,l')*)
      | File elt -> File (f elt)
    )
;;
(*let rec map_fd (f: 'a -> 'b directory) (dir: 'a directory) : 'b directory =*)
let rec map_fd ?(name=None) ~f dir =
    ( match dir with 
      | Directory (name',l) -> 
          let l' = List.map (fun dir -> map_fd ~f dir) l in
	      make (match name with Some s -> s | None -> name') l'
(*          Directory (name,l')*)
      | File elt -> f elt
    )
;;

let merge_files' ~(g: 'a list -> 'a directory list -> 'b directory list)
  (*~(f: 'a list -> 'b)*)
   (dir: 'a directory) : 'b directory =
    ( match dir with
      | Directory (name,l) ->
            let lst_files' = List.fold_left (fun acc elt ->
                match elt with File file -> file::acc | _ -> acc)
                [] l in
            let lst_files = List.rev lst_files' in
            let lst_dir = List.filter (
                function Directory _ -> true | _ -> false)
                l in
            Directory (name, 
               g lst_files lst_dir
(*               (File (f lst_files))::(List.map (merge_files ~f) lst_dir)*)
            )
      | File file -> 
              failwith "merge_files'"
(*              let file' = File (f [file]) in
              Directory ("merged", [file'])*)
(*              of_files "merged" [f [file]]*)
    )
;;

let rec merge_files ~f dir = merge_files' dir 
        ~g:(fun files dirs -> 
            (File (f files))::(List.map (merge_files ~f) dirs))

let rec merge_files_rm_empty ~f dir = merge_files' dir 
        ~g:(fun files dirs -> 
            let l = List.map (merge_files_rm_empty ~f) dirs in
            if files = [] 
            then l
            else (File (f files))::l )

let rec rearrange_files ~(f: 'a list -> 'a list) (dir: 'a directory) : 'a directory =
    ( match dir with
      | Directory (name,l) ->
            let lst_files' = List.fold_left (fun acc elt ->
                match elt with File file -> file::acc | _ -> acc)
                [] l in
            let lst_files = List.rev lst_files' in
            let lst_dir = List.filter (
                function Directory _ -> true | _ -> false)
                l in
            let new_files = List.map (fun file -> File file) (f lst_files) in
            Directory (name, 
               new_files@(List.map (rearrange_files ~f) lst_dir)
            )
      | File file -> File file
    )
;;

let glue_elt dir1 dir2 = 
    ( match dir1, dir2 with
      | Directory(name1,l1), Directory(name2,l2) ->
              Directory(name1, l1@l2)
      | Directory(name1,l1), File f2 ->
              Directory(name1, l1@[File f2])
      | File f1, Directory(name2,l2) ->
              Directory(name2, (File f1)::l2)
      | File f1, File f2 ->
              Directory("join", [File f1; File f2])
    )
;;

let is_file = function File _ -> true | _ -> false;;
let is_directory = function Directory _ -> true | _ -> false;;

let rec join = 
  ( function
    | Directory (name,l) -> Directory (name, join_list l)
    | File f -> File f)
and join_list =
  ( function 
    | b1::bs ->
        ( match b1 with
          | File f -> (File f)::(join_list bs)
          | Directory (name,_) ->
               let cond = ( function 
                   Directory (name',_) -> name = name'
                 | File _ -> false
               ) in
               let same_name = List.filter cond bs and
               other_name = List.filter (fun elt -> not (cond elt)) bs in
               let new_b1 = List.fold_left glue_elt b1 same_name in
               (join new_b1)::(join_list other_name)
        )
    | [] -> []
  )
;;

let join2 dir1 dir2 = join (glue_elt dir1 dir2)

(*let join = glue_elt*)
(*
let rec join dir1 dir2 = 
    ( match dir1, dir2 with
      | Directory(name1,l1), Directory(name2,l2) ->
              if name1 = name2
              then (
                  let new_l = List.fold_left (fun acc elt1 ->
*)


let rec iter_files ~(f: 'a -> unit) (dir: 'a directory) : unit =
    match dir with
      | Directory (_,l) -> List.iter (iter_files ~f) l
      | File elt -> f elt
;;

(** [Iterators on two directories] *)
(** Directories' structure and names should strictly coincide! 
 Otherwise exception 'Error_iter2' txt will be raised *)
exception Error_iter2 of string

let rec iter2_files ~(f: 'a -> 'b -> unit) 
        (dir1: 'a directory)
        (dir2: 'b directory) : unit =
    match dir1, dir2 with
      | Directory (name1,l1), Directory (name2,l2) -> 
        if name1 = name2 then (
            try
                List.iter2 (iter2_files ~f) l1 l2
            with Invalid_argument s -> 
            raise (Error_iter2 (
                "iter2_files: Invalid_argument '"^s^"' "^
                "(lists have different lengths?) "^
                "in directory '"^name1^"'"))
        ) else (
            raise (Error_iter2 (
                "directories names mismatch: "^
                "'"^name1^"', '"^name2^"'"))
        )
      | File elt1, File elt2 -> f elt1 elt2
      | Directory (name,_), _ | _, Directory (name,_) ->
        raise  (Error_iter2 (
            "directory structure do not coincide near "^
            "directory '"^name^"'"))
;;

let rec map2 ~f dir1 dir2 =
    match dir1, dir2 with
      | Directory (name1,l1), Directory (name2,l2) -> 
        if name1 = name2 then (
            try
                Directory (name1, List.map2 (map2 ~f) l1 l2)
            with Invalid_argument s -> 
            raise (Error_iter2 (
                "map2_files: Invalid_argument '"^s^"' "^
                "(lists have different lengths?) "^
                "in directory '"^name1^"'"))
        ) else (
            raise (Error_iter2 (
                "directories names mismatch: "^
                "'"^name1^"', '"^name2^"'"))
        )
      | File elt1, File elt2 -> File (f elt1 elt2)
      | Directory (name,_), _ | _, Directory (name,_) ->
        raise  (Error_iter2 (
            "directory structure do not coincide near "^
            "directory '"^name^"'"))
;;


let rec iter_path_files ?(current_path="") ~(f: path:string -> 'a -> unit) (dir: 'a directory) : unit =
    match dir with
      | Directory (name,l) -> 
              let path = current_path^"/"^name in
              List.iter (iter_path_files ~current_path:path ~f) l
      | File elt -> f ~path:current_path elt
;;

let fold_files ~(f: 'b -> 'a -> 'b) ~(init: 'b) (dir: 'a directory) : 'b =
    let result = ref init in
    iter_files ~f:(fun file -> result := f !result file) dir;
    !result
;;

let fold2_files ~(f: 'b -> 'a -> 'c -> 'b) ~(init: 'b) 
        (dir1: 'a directory) 
        (dir2: 'c directory) : 'b =
    let result = ref init in
    iter2_files ~f:(fun file1 file2 -> result := f !result file1 file2) 
        dir1 dir2;
    !result
;;

let fold_path_files ~f ~init dir =
    let result = ref init in
    iter_path_files ~f:(fun ~path file -> result := f ~path !result file) dir;
    !result
;;

module Path = struct
  let to_list path =
      let path' = Str.global_replace (Str.regexp "/+") "/" path in
      Str.split (Str.regexp "/") path' 

  let tl path =
      let l = List.tl (to_list path) in
      List.fold_left (fun acc elt -> acc^"/"^elt) "" l
end

module Simple_make ( Content: sig type t end) =
struct
  type content = Content.t
  type t = Content.t directory

  let make name l : t = 
      Directory (name,l)
  let of_files name fl : t =
      let l' = List.map (fun elt -> File elt) fl in
      Directory (name, l')
  let map f (d : t) = map
  let iter_files f (d : t) = iter_files f d
  let fold_files f (init: content) (d : t) = fold_files f init d
end

(* TODO: test on correctness during creation *)

exception Find_error of string
let fail s = raise (Find_error s)

module Find_print_make ( Content: sig 
    type t
    val name_of:  t -> string
(*    val equiv: string -> t -> bool*)
end) =
struct
  include Simple_make(Content)

  let name_of (dir : t) = ( match dir with
      Directory (name,_) -> name
    | File file -> Content.name_of file)

(*  let find_sub str = (function
      Directory (_, l) ->
        ( try
           Some (List.find (fun elt -> name elt = str) l)
           with Not_found -> None)
    | _ -> None
  )*)
(*  let rec find_pl path_list (dir : t) = 
      ( match path_list with
        | [] -> None
        | name::pl -> 
              ( match find_sub name dir with
                | Some dir -> find_pl pl dir
                | None -> None
              )
      )*)
(*
  let find_file path dir =
      let path_list = Str.split (Str.regexp "/") path in
      ( match find_pl path_list dir with
        | Some (File file) -> Some file
        | _ -> None
      )
*)
  let rec get_elt (elt: t) ?(current_path="") path = 
    let wrong_demand name i = fail (
          "get_elt: "^(string_of_int i)^". demand get '"^name^"' from "^
          (type_of elt)^" '"^(name_of elt)^"'.") in
    (match path with
      | [] -> fail "get_elt: empty demand"
      | name::[] -> 
          if name = name_of elt 
          then elt
          else wrong_demand name 1
      | name::next_name::path' ->
        (match elt with
          | Directory (current_name,l) -> 
            if name <> current_name
            then wrong_demand name 2
            else (
              let next_elt = (try
                List.find (fun elt -> name_of elt = next_name) l
                with Not_found -> (
                    let elt_lst = List.fold_left (fun acc elt ->
                        acc^" "^(name_of elt)^";") "" l in
                    fail ("get_elt: no element named '"^name^
                    "' is available among ["^elt_lst^"]")
                )) in
              get_elt next_elt (next_name::path') 
                  ~current_path:(current_path^"/"^name)
            )
          | File file ->
             wrong_demand name 3
(*            fail ("element '"^(current_path^"' is not directory")*)
          )
  )
  let find_elt ~path (dir: t) =
      let path_list = Path.to_list path in
      get_elt dir path_list

  let find_file ~path (dir: t) =
      ( match find_elt ~path dir with
        | File file -> file
        | _ -> raise (Find_error (path^"is directory!"))
      )
(*
  let find_sub str = (function
      Directory (_, l) ->
          (try 
           List.find (fun elt -> name elt = str) l
          with Not_found -> (
              let elt_lst = List.fold_left (fun acc elt ->
                  acc^" "^(name elt)) "" l in
              fail ("no element '"^str^"', available: "^elt_lst)
          ))
    | _ -> raise (Find_error "not dir")
  )
  let rec find_pl path_list (dir : t) = 
      ( match path_list with
        | [] -> raise (Find_error "empty demand")
        | name::pl -> 
              let r = find_sub name dir in
              find_pl pl r
      )
  let find_file path (dir: t) =
      let path_list = Str.split (Str.regexp "/") path in
      ( match find_pl path_list dir with
        | File file -> file
        | _ -> raise (Find_error (path^"is directory!"))
      )
*)
  let print_shift n = print_string (String.make (3*n) ' ')

  let rec print ?(shift=0) (dir: t) =
      ( match dir with
        | Directory (name,l) ->
            print_shift shift; 
            print_endline (name^"/");
            List.iter (print ~shift:(shift+1)) l
        | File file ->
            print_shift shift; 
            print_endline (Content.name_of file)
      )
end


module Make_state(T: sig
    type file
    type value_t
    val init: file -> value_t 
    type data
    val get_value: file -> data -> value_t 
end) =
struct

  type value_t = T.value_t
  type t = (value_t ref) directory 

  let init (dir: T.file directory) : t =
      map (fun file -> ref (T.init file)) dir
  ;;

  let set (state: t) ~dir (r: T.data) =
      iter2_files state dir ~f:(fun a b ->
          a := T.get_value b r)
  ;;

  let collect2 ~(state: t) (dir: 'a directory) ~init
        ~(f: value_t ref -> 'a -> 'b) : 'b list =
      fold2_files ~init
        state dir
        ~f:(fun acc a b -> (f a b) :: acc)
  ;;
end

